import tensorflow as tf
import tensorflow.contrib.slim as slim
from data_util.streamer import Oulu_NPU, OneLabelData, HomeMadeData
from data_util.parser import *
import os
import numpy as np
import tensorflow_hub as hub
import sys
from data_util.data_ops import *
from tensorflow.python.ops import init_ops

train_data_ops = [ hotel_train_pos_ops , hotel_train_neg_ops ]
test_data_ops  = [ hotel_test_pos_ops , hotel_test_neg_pos ]
test_writers   = [ hotel_pos_writer , hotel_neg_writer ]

model_dir = "/home/jh/working_pros/binary_classifier/hard2Train_model"

if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "3"

    tf.logging.set_verbosity( tf.logging.INFO )

    sess = tf.Session()
    graph = tf.get_default_graph()

    saver = tf.train.import_meta_graph( os.path.join( model_dir , "model.ckpt.meta") )
    saver.restore( sess , tf.train.latest_checkpoint( model_dir ) )

    imgs = graph.get_tensor_by_name( "module/hub_input/images:0" )
    label_placeholder = graph.get_tensor_by_name( "label:0" )

    acc = graph.get_tensor_by_name( "acc:0" )
    loss = graph.get_tensor_by_name( "loss/value:0" )
    acc_summary = graph.get_tensor_by_name( "accuracy:0" )
    loss_summary = graph.get_tensor_by_name( "loss_1:0" )

    trainable_list = tf.trainable_variables()
    for t in trainable_list:
        print( t )

    train_op = graph.get_operation_by_name( "Adam" )

    summary_merged = tf.summary.merge( [ loss_summary , acc_summary ] )
    train_writer = tf.summary.FileWriter( './tflog/train' , graph = tf.get_default_graph() )

    saver = tf.train.Saver()

    for i in range( 10000 ):
        t_I = []
        t_L = []

        for op in train_data_ops:
            tii, tll = sess.run( op )
            t_I.append( tii )
            t_L.append( tll )

        t_I = np.concatenate( t_I , axis = 0 )
        t_L = np.concatenate( t_L , axis = 0 )

        _ , ACC , LOSS, SUMMARY = sess.run( \
                [train_op , acc , loss , summary_merged ] , \
                feed_dict = { imgs: t_I , \
                label_placeholder : t_L } )
        
        train_writer.add_summary( SUMMARY , i )
        print( "iter = %d , loss = %f "  %( i , LOSS ) )

        if i%200 == 0:
            save_path = saver.save(sess, "./tflog/tmp/model.ckpt")
            print( "Model saved in path: %s" % save_path )

        if i%50 == 0:
            for op, writer in zip( test_data_ops , test_writers ):
                t_I , t_L = sess.run( op )

                ACC , LOSS, SUMMARY = sess.run( \
                        [ acc , loss , summary_merged ] , \
                        feed_dict = { imgs: t_I , \
                        label_placeholder : t_L } )

                writer.add_summary( SUMMARY , i )
                print( "accuracy accuracy accuracy    =     %f" % ACC )
